package com.yuanxiao.common.core.utils;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;

public class TreeUtils {

    /**
     * 构建 N 层树结构（迭代方式，支持自定义根节点父 ID）
     *
     * @param list           扁平数据列表
     * @param idGetter       获取节点 ID 的方法
     * @param parentIdGetter 获取父节点 ID 的方法
     * @param childrenSetter 设置子节点的方法
     * @param rootParentId   根节点的父 ID（如 0L、-1L、null）
     * @param <T>            泛型类型
     * @return 树形结构根节点列表
     */
    public static <T> List<T> buildTree(
            List<T> list,
            Function<T, Long> idGetter,
            Function<T, Long> parentIdGetter,
            BiConsumer<T, List<T>> childrenSetter,
            Long rootParentId) {

        if (list == null || list.isEmpty()) {
            return Collections.emptyList();
        }

        Map<Long, T> map = new HashMap<>();
        Map<Long, List<T>> parentChildrenMap = new HashMap<>();

        // 构建 ID -> 节点 映射
        list.forEach(item -> map.put(idGetter.apply(item), item));

        // 构建 父ID -> 子节点列表 映射
        list.forEach(item -> {
            Long parentId = parentIdGetter.apply(item);
            // 统一处理 null，转换为 rootParentId（如果需要）
            if (parentId == null && (rootParentId != null)) {
                parentId = rootParentId;
            }
            parentChildrenMap.computeIfAbsent(parentId, k -> new ArrayList<>()).add(item);
        });

        // 获取根节点列表
        List<T> rootNodes = parentChildrenMap.getOrDefault(rootParentId, new ArrayList<>());
        Queue<T> queue = new LinkedList<>(rootNodes);

        // 迭代构建树
        while (!queue.isEmpty()) {
            T current = queue.poll();
            Long currentId = idGetter.apply(current);
            List<T> children = parentChildrenMap.getOrDefault(currentId, Collections.emptyList());
            childrenSetter.accept(current, children);
            queue.addAll(children);
        }

        return rootNodes;
    }
}
